from __future__ import with_statement

import unittest

import time, timeit
import math, logging

import numpy, scipy.constants, scipy.integrate
import numpy.testing

import wlc

logger = logging.getLogger("wlc_test")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

class WlcTest(unittest.TestCase):

    def test_f_inverse(self):
        kb_t_over_p = 2.0
        v_over_lc = 5.0
        k_lc = 4.0
        y = 3.0
        t = wlc.f_inverse(kb_t_over_p, k_lc, v_over_lc, y)
        x_over_lc = v_over_lc * t - y / k_lc
        actual = wlc.wlc(kb_t_over_p, x_over_lc)
        numpy.testing.assert_almost_equal(y, actual)

    def test_f_inverse_derived(self):
        kb_t_over_p = 2.0
        v_over_lc = 5.0
        k_lc = 4.0
        y1 = 3.0
        t1 = wlc.f_inverse(kb_t_over_p, k_lc, v_over_lc, y1)
        y2 = y1 + 2.0 ** -20
        t2 = wlc.f_inverse(kb_t_over_p, k_lc, v_over_lc, y2)
        expected = (t2 - t1) / (y2 - y1)
        actual = wlc.f_inverse_derived(kb_t_over_p, k_lc, v_over_lc, y1)
        numpy.testing.assert_allclose(actual, expected)
        actual2 = wlc.f_inverse_derived_approximation(
            kb_t_over_p, k_lc, v_over_lc, y1)
        numpy.testing.assert_allclose(actual2, expected, rtol=0.01)

    def test_f_derived(self):
        kb_t_over_p = 2.0
        v_over_lc = 5.0
        k_lc = 4.0
        t1 = 3.0
        x_over_lc1 = wlc.find_x_over_lc_for_tug_of_war(
            kb_t_over_p, k_lc, v_over_lc * t1)
        f1 = wlc.wlc(kb_t_over_p, x_over_lc1)
        t2 = t1 + 2.0 ** -20
        x_over_lc2 = wlc.find_x_over_lc_for_tug_of_war(
            kb_t_over_p, k_lc, v_over_lc * t2)
        f2 = wlc.wlc(kb_t_over_p, x_over_lc2)
        expected = (f2 - f1) / (t2 - t1)
        actual = wlc.f_derived(kb_t_over_p, k_lc, v_over_lc, t1)
        numpy.testing.assert_allclose(actual, expected)

    def test_distance_over_extension(self):
        kb_t_over_p = scipy.constants.k * 300 / wlc.default_p
        for x_over_lc in numpy.arange(0, 1, 0.01):
            f = wlc.wlc(kb_t_over_p, x_over_lc)
            logger.debug("x_over_lc=%s f[pN]=%s", x_over_lc, f * 1e12)
            actual = wlc.compute_x_over_lc(kb_t_over_p, f)
            logger.debug("actual=%r", actual)
            numpy.testing.assert_approx_equal(actual, x_over_lc)

    def test_find_x_for_tug_of_war(self):
        kb_t_over_p = scipy.constants.k * 300 / wlc.default_p
        lc = 18e-9
        k = 0.06
        for z in numpy.arange(1e-9, 10e-9, 1e-10):
            x = wlc.find_x_for_tug_of_war(kb_t_over_p, lc, k, z)
            # logger.info("x=%r z=%r", wlc.wlc(kb_t_over_p, x / lc), (z - x) * k)
            self.assertTrue(x < z)
            numpy.testing.assert_approx_equal(
                wlc.wlc(kb_t_over_p, x / lc), (z - x) * k)

    def test_radial_probability_density(self):
        l = 18e-9
        actual, _ = scipy.integrate.quad(
            lambda r: wlc.radial_probability_density(wlc.default_p / l, r), 0, 1.0)
        numpy.testing.assert_allclose(actual, 1, rtol=0.02)

    def test_log_radial_probability_density(self):
        l = numpy.longfloat(18e-9)
        for r in numpy.linspace(0, 1, 100):
            r = numpy.longfloat(r)
            actual = wlc.log_radial_probability_density(wlc.default_p / l, r)
            p = wlc.radial_probability_density(wlc.default_p / l, r)
            expected = numpy.log(p) if p else -numpy.inf
            logger.debug("r=%s actual=%s expected=%s", r, actual, expected)
            numpy.testing.assert_allclose(actual, expected)

def find_x_for_tug_of_war_for_profiling():
    kb_t_over_p = scipy.constants.k * 300 / wlc.default_p
    lc = 18e-9
    k = 0.06
    xs = []
    for z in numpy.arange(1e-9, 10e-9, 1e-10):
        x = wlc.find_x_for_tug_of_war(kb_t_over_p, lc, k, z)
        xs.append(x)
    return xs

if __name__ == "__main__":
    numpy.seterr(divide="raise", invalid="raise")
    suite = unittest.TestLoader().loadTestsFromTestCase(WlcTest)
    unittest.TextTestRunner(verbosity=2).run(suite)
    # unittest.main() 
